# # backend/main.py
# from fastapi import FastAPI, HTTPException
# from fastapi.middleware.cors import CORSMiddleware
# from pydantic import BaseModel
# from typing import List
# import uvicorn
# from rag import costomed_retriever, vs, get_answer

# app = FastAPI(title="问答系统API")

# # 允许跨域请求
# app.add_middleware(
#     CORSMiddleware,
#     allow_origins=["*"],
#     allow_methods=["*"],
#     allow_headers=["*"],
# )

# class QuestionRequest(BaseModel):
#     question: str
#     k: int = 2
#     score_threshold: float = 0.4

# class DocumentResponse(BaseModel):
#     content: str

# class AnswerResponse(BaseModel):
#     answer: str
#     documents: List[DocumentResponse]

# @app.post("/query", response_model=AnswerResponse)
# async def query_policies(request: QuestionRequest):
#     """查询接口"""
#     try:
#         # 获取回答和相关文档
#         answer = get_answer(request.question)
#         docs = costomed_retriever(request.question, vs, k=request.k, score_threshold=request.score_threshold)
        
#         # 格式化文档
#         formatted_docs = [DocumentResponse(content=doc.page_content) for doc in docs]
        
#         return AnswerResponse(answer=answer, documents=formatted_docs)
    
#     except Exception as e:
#         raise HTTPException(status_code=500, detail=str(e))

# if __name__ == "__main__":
#     uvicorn.run(app, host="127.0.0.1", port=8000)


import os
from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from typing import List, Optional
import uvicorn
# from rag import costomed_retriever, vs, get_answer
from rag import PolicyRAGSystem
from langchain.vectorstores import FAISS
from sqlalchemy import create_engine, text
import datetime
from search_policy_by_tag import search_policy_by_tags


# --- 数据库配置 ---
DB_CONFIG = {
    'host': 'localhost',
    'port': '5432',
    'database': 'Policy_DB',
    'user': 'postgres',
    'password': 'postgres'
}

# 创建数据库引擎
db_url = f"postgresql://{DB_CONFIG['user']}:{DB_CONFIG['password']}@{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}"
engine = create_engine(db_url)

# --- 配置部分 ---
class Config:
    API_TITLE = "政策问答系统API"
    API_DESCRIPTION = "基于RAG架构的政策咨询系统后端服务"
    API_VERSION = "1.0.0"
    HOST = os.getenv("HOST", "0.0.0.0")
    PORT = int(os.getenv("PORT", 8000))
    RELOAD = os.getenv("RELOAD", "True")
    LOG_LEVEL = os.getenv("LOG_LEVEL", "info")

# --- 数据模型 ---
class QueryRequest(BaseModel):
    question: str
    k: Optional[int] = 2
    score_threshold: Optional[float] = 0.4

class TagsRequest(BaseModel):
    tags: list

class DocumentResponse(BaseModel):
    content: str
    score: Optional[float] = None

class AnswerResponse(BaseModel):
    answer: str

class PolicyResponse(BaseModel):
    answer: list[dict]

class PolicyCountResponse(BaseModel):
    status: str
    data: dict

class HealthCheckResponse(BaseModel):
    status: str
    model_loaded: bool
    service_version: str

class PolicySummaryResponse(BaseModel):
    policy_total_count: int
    today_update_count: int

# --- 应用初始化 ---
app = FastAPI(
    title=Config.API_TITLE,
    description=Config.API_DESCRIPTION,
    version=Config.API_VERSION,
    docs_url="/api/docs",
    redoc_url="/api/redoc"
)

# CORS中间件配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 生产环境应限制为具体域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# --- 中间件 ---
@app.middleware("http")
async def add_process_time_header(request: Request, call_next):
    """添加请求处理时间头"""
    from time import time
    start_time = time()
    response = await call_next(request)
    process_time = time() - start_time
    response.headers["X-Process-Time"] = str(process_time)
    return response

# --- 路由 ---
@app.post("/api/query", 
          response_model=AnswerResponse,
          summary="政策问答查询",
          description="提交问题获取政策解答和相关文档")
async def query_policies(request: QueryRequest):
    """
    示例请求体：
    {
        "question": "高新技术企业有哪些税收优惠",
        "k": 3,
        "score_threshold": 0.4
    }
    """
    try:
        # 获取回答和相关文档
        policy_rag_system = PolicyRAGSystem()
        dir = "total_db"
        policy_rag_system.vector_store = FAISS.load_local(
            dir, 
            policy_rag_system.embedding_model,
            allow_dangerous_deserialization=True
        )
        answer = policy_rag_system.generate_answer(request.question)
        # docs = costomed_retriever(
        #     request.question, 
        #     vs, 
        #     k=request.k, 
        #     score_threshold=request.score_threshold
        # )
        
        # # 格式化文档响应
        # formatted_docs = [
        #     DocumentResponse(
        #         content=doc.page_content,
        #         score=score
        #     ) for doc, score in [
        #         (doc, score) for doc, score in vs.similarity_search_with_score(
        #             request.question, 
        #             k=request.k
        #         ) if score >= request.score_threshold
        #     ]
        # ]
        
        return AnswerResponse(
            answer=answer,
            # documents=formatted_docs
        )
        
    except Exception as e:
        print(f"查询处理失败: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"查询处理失败: {str(e)}"
        )

@app.get("/api/health",
         response_model=HealthCheckResponse,
         summary="服务健康检查")
async def health_check():
    """服务状态检查端点"""
    return {
        "status": "healthy",
        "model_loaded": True,  # 可替换为实际模型加载检查
        "service_version": Config.API_VERSION,
        "db_connected": await check_db_connection()
    }

# --- 数据库健康检查 ---
async def check_db_connection():
    try:
        with engine.connect() as conn:
            conn.execute(text("SELECT 1"))
        return True
    except Exception:
        return False
    
@app.get("/api/policy/get_total_count",
         response_model=PolicyCountResponse,
         summary="获取政策总数")
async def get_policy_count():
    """获取政策库中的政策总数"""
    try:
        with engine.connect() as conn:
            result = conn.execute(text("SELECT COUNT(*) FROM policy_content_table"))
            count = result.scalar()
        
        return {
            "status": "success",
            "data": {"policy_total_count": count}
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    

    
@app.get("/api/policy/get_total_language_count",
         summary="获取政策覆盖语言总数")
async def get_total_language_count():
    """获取政策库中的政策覆盖语言总数"""
    try:
        with engine.connect() as conn:
            result = conn.execute(text("SELECT COUNT(DISTINCT(language)) FROM policy_content_table"))
            count = result.scalar()
        
        return {
            "status": "success",
            "data": {"policy_total_language_count": count}
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    
@app.get("/api/policy/get_total_crossborder_policy_count",
         summary="获取跨境政策数")
async def get_total_crossborder_policy_count():
    """获取政策库中的跨境政策总数"""
    try:
        with engine.connect() as conn:
            result = conn.execute(text("SELECT COUNT(*) FROM policy_content_table WHERE is_crossborder_policy"))
            count = result.scalar()
        
        return {
            "status": "success",
            "data": {"policy_total_crossborder_policy_count": count}
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/policy/get_today_update_count",
         response_model=PolicyCountResponse,
         summary="获取今日更新政策数")
async def get_today_update_count():
    """获取今日更新的政策数量"""
    try:
        today = datetime.datetime.now().strftime('%Y-%m-%d')
        with engine.connect() as conn:
            result = conn.execute(
                text("SELECT COUNT(*) FROM policy_content_table WHERE DATE(publish_time) = :today"),
                {'today': today}
            )
            count = result.scalar()
        
        return {
            "status": "success",
            "data": {"today_update_count": count}
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    
@app.get("/api/policy/get_total_province_count",
         summary="获取政策覆盖的省份总数")
async def get_total_province_count():
    """获取政策库中的政策覆盖的省份总数"""
    try:
        with engine.connect() as conn:
            result = conn.execute(text("SELECT COUNT(DISTINCT(address_2)) FROM policy_content_table"))
            count = result.scalar()
        
        return {
            "status": "success",
            "data": {"policy_total_province_count": count}
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    
@app.get("/api/policy/get_total_city_count",
         summary="获取政策覆盖的城市总数")
async def get_total_city_count():
    """获取政策库中的政策覆盖的城市总数"""
    try:
        with engine.connect() as conn:
            result = conn.execute(text("SELECT COUNT(DISTINCT(address_3)) FROM policy_content_table"))
            count = result.scalar()
        
        return {
            "status": "success",
            "data": {"policy_total_city_count": count}
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/policy/get_total_department_count",
         summary="获取政策覆盖的发文单位总数")
async def get_total_department_count():
    """获取政策库中的政策覆盖的发文单位总数"""
    try:
        with engine.connect() as conn:
            result = conn.execute(text("SELECT COUNT(DISTINCT(department)) FROM policy_content_table"))
            count = result.scalar()
        
        return {
            "status": "success",
            "data": {"policy_total_department_count": count}
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/policy/summary-stats",
         response_model=None,
         summary="获取政策摘要统计")
async def get_policy_summary_stats():
    """一次性获取所有关键统计信息"""
    try:
        with engine.connect() as conn:
            # 获取政策总数
            total_result = conn.execute(text("SELECT COUNT(*) FROM policy_content_table"))
            total_count = total_result.scalar()
            
            # 获取今日更新数
            today = datetime.datetime.now().strftime('%Y-%m-%d')
            today_result = conn.execute(
                text("SELECT COUNT(*) FROM policy_content_table WHERE DATE(publish_time) = :today"),
                {'today': today}
            )
            today_count = today_result.scalar()

            # 获取覆盖的语言数量
            total_language_result = conn.execute(text("SELECT COUNT(DISTINCT language) FROM policy_content_table"))
            total_language_count = total_language_result.scalar()

            # 可以添加其他统计查询...
            
        return {
            "status": "success",
            "data": {
                "policy_total_count": total_count,
                "today_update_count": today_count,
                "total_language_count": total_language_count
                # 其他统计数据...
            }
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/policy/search_by_tags",
          summary = "根据标签筛选政策")
async def search_by_tags(request: TagsRequest):
    try:
        res = search_policy_by_tags(request.tags)
        return PolicyResponse(answer=res)
            
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    



# --- 错误处理 ---
@app.exception_handler(HTTPException)
async def http_exception_handler(request, exc):
    return JSONResponse(
        status_code=exc.status_code,
        content={"detail": exc.detail},
    )

# --- 启动配置 ---
if __name__ == "__main__":
    uvicorn.run(
        "backend:app",
        host=Config.HOST,
        port=Config.PORT,
        reload=Config.RELOAD,
        log_level=Config.LOG_LEVEL,
        workers=int(os.getenv("WORKERS", 1))
    )